﻿#include "Atlas.h"

// lib
#include <stb_image.h>
#include <json.hpp>

// std
#include <filesystem>
#include <fstream>

Atlas::Atlas(uint32_t target, uint32_t internalFormat, uint32_t format, uint32_t type, bool generateMipMap, int32_t maxLod)
{
  Generate(target, internalFormat, format, type, generateMipMap, maxLod);
}

Atlas::Atlas(const std::string &directoryPath, int atlasWidth, int atlasHeight, uint32_t target, uint32_t internalFormat, uint32_t format, uint32_t type, bool generateMipMap, int32_t maxLod)
    : Atlas(target, internalFormat, format, type, generateMipMap, maxLod)
{
  CreateFromDirectory(directoryPath, atlasWidth, atlasHeight);
}

Atlas::Atlas(const std::string &jsonPath)
{
  CreateFromJson(jsonPath);
}

void Atlas::CreateFromDirectory(const std::string &directoryPath, int atlasWidth, int atlasHeight)
{
  LOG_INFO("Creating Atlas from: {}", directoryPath);

  int nextX = 0;
  int nextY = 0;

  m_height = atlasHeight;
  m_width = atlasWidth;
  m_channels = 4;

  int atlasSize = atlasWidth * atlasHeight * 4;
  uint8_t *atlasPixels = new uint8_t[atlasSize];
  memset(atlasPixels, 0, atlasSize);

  // 遍历目录以及子目录下的所有文件，stbi_load加载，并复制到一张大纹理中，编排为图集
  for (const auto &entry : std::filesystem::recursive_directory_iterator(directoryPath))
  {
    auto ext = entry.path().extension();
    if (!entry.is_regular_file() || ext != ".png")
      continue;

    auto extSize = ext.string().size();

    std::string filepath = entry.path().generic_string();

    int texWidth, texHeight, texChannels;

    // 加载图片文件并获取像素数据
    stbi_set_flip_vertically_on_load(true);
    stbi_uc *pixels =
        stbi_load(
            filepath.c_str(), &texWidth, &texHeight,
            &texChannels,
            STBI_rgb_alpha);

    if (!pixels)
      continue;

    int row = texHeight / TileSize;
    int col = texWidth / TileSize;
    int atlasRow = atlasHeight / TileSize;
    int atlasCol = atlasWidth / TileSize;

    const std::string texName = m_prefix +
                                filepath.substr(
                                    directoryPath.size() + 1,
                                    filepath.size() - directoryPath.size() - extSize - 1);

    for (int i = 0; i < row; i++)
    {
      for (int j = 0; j < col; j++)
      {
        // 找到一个空位置放置图片
        int x = nextX, y = nextY;
        if (y >= atlasHeight)
        {
          LOG_ERROR("Atlas is full!");
          return;
        }

        // 复制图片到大纹理中
        for (int py = 0; py < TileSize; ++py)
        {
          for (int px = 0; px < TileSize; ++px)
          {
            const int offset = (py * TileSize + px) * 4;
            // 计算源像素坐标（当前tile在原始图片中的位置）
            const int srcX = j * TileSize;
            const int srcY = i * TileSize;

            // 源像素索引（原始图片坐标系）
            const int srcIndex = ((srcY + py) * texWidth + (srcX + px)) * 4;

            // 目标像素索引（图集坐标系）
            const int dstIndex = ((y + py) * atlasWidth + (x + px)) * 4;

            // 复制所有通道（RGBA）
            atlasPixels[dstIndex + 0] = pixels[srcIndex + 0]; // R
            atlasPixels[dstIndex + 1] = pixels[srcIndex + 1]; // G
            atlasPixels[dstIndex + 2] = pixels[srcIndex + 2]; // B
            atlasPixels[dstIndex + 3] = pixels[srcIndex + 3]; // A
          }
        }

        if (!m_textureUVMap.contains(texName))
          m_textureUVMap[texName] = {{x / TileSize, y / TileSize}, 1};
        else
          m_textureUVMap[texName].second++;

        // 更新下一个位置
        nextX += TileSize;
        if (nextX >= atlasWidth)
        {
          nextX = 0;
          nextY += TileSize;
        }
      }
    }

    assert(m_textureUVMap.contains(texName));
    auto &[uv, animSize] = m_textureUVMap[texName];
    LOG_TRACE("Texture: {} UV: ({}, {}), AnimSize: {}", texName, uv.x / TileSize, uv.y / TileSize, animSize);

    stbi_image_free(pixels);
  }

  Bind();

  GLenum format = m_channels == 4 ? GL_RGBA : GL_RGB;
  glTexImage2D(m_target, 0, format, atlasWidth, atlasHeight, 0, format, GL_UNSIGNED_BYTE, atlasPixels);

  if (m_generateMipMap)
  {
    glGenerateMipmap(m_target);
  }

  delete[] atlasPixels;
}

void Atlas::CreateFromJson(const std::string &jsonPath)
{
  std::ifstream file(jsonPath);
  if (!file.is_open())
  {
    LOG_ERROR("Failed to open JSON file: {}", jsonPath);
    return;
  }

  nlohmann::json jsonData;
  file >> jsonData;

  const std::string directoryPath = jsonData["directory"];
  const int atlasWidth = jsonData["width"];
  const int atlasHeight = jsonData["height"];
  m_prefix = jsonData.contains("prefix") ? jsonData["prefix"] : "";

  int maxLod = 1;
  if (jsonData.contains("maxLod")) 
    maxLod = jsonData["maxLod"];

  Generate(GL_TEXTURE_2D, GL_RGBA, GL_RGBA, GL_UNSIGNED_BYTE, true, maxLod);

  CreateFromDirectory(directoryPath, atlasWidth, atlasHeight);
}

std::pair<glm::ivec2, int> Atlas::GetTextureIndex(const std::string &textureName) const
{
  if (m_textureUVMap.contains(textureName))
  {
    return m_textureUVMap.at(textureName);
  }
  return {{0, 0}, 1};
}

bool Atlas::FindEmpityPosition(uint8_t *atlasPixels, int atlasWidth, int atlasHeight, int texWidth, int texHeight, int &outX, int &outY)
{
  // 由于贴图长宽基本是16的倍数，所以可以从16的倍数开始遍历，减少遍历次数
  bool foundPosition = false;
  for (int x = 0; x < atlasWidth - texWidth + 1; x += 16)
  {
    for (int y = 0; y < atlasHeight - texHeight + 1; y += 16)
    {
      // 检查当前位置是否可以放置图片
      bool canPlace = true;
      for (int py = 0; py < texHeight; ++py)
      {
        for (int px = 0; px < texWidth; ++px)
        {
          const int srcIndex = (py * texWidth + px) * 4;
          const int dstIndex = ((y + py) * atlasWidth + (x + px)) * 4;
          if (atlasPixels[dstIndex + 3] > 0) // 检查 alpha 通道是否为非零
          {
            canPlace = false;
            break;
          }
        }
        if (!canPlace)
          break;
      }
      if (canPlace)
      {
        foundPosition = true;
        outX = x;
        outY = y;
        break;
      }
    }
    if (foundPosition)
      break;
  }
  return foundPosition;
}
